import os
from concurrent.futures import ThreadPoolExecutor
from PIL import Image
import re

def crop_and_resize(image_path, label_path, output_folder, scale_factor=1.5, quality=85):
    """
    裁剪图片并放大每个bbox，然后保存图片
    :param image_path: 图片的路径
    :param label_path: 对应的YOLO格式的标注文件路径
    :param output_folder: 输出文件夹
    :param scale_factor: 放大倍数
    :param quality: 保存图像的质量
    """
    if not os.path.exists(image_path) or not os.path.exists(label_path):
        print(f"File not found: {image_path} or {label_path}")
        return
    
    # 打开图片
    img = Image.open(image_path)
    img_width, img_height = img.size

    # 读取YOLO标签文件
    with open(label_path, "r") as label_file:
        lines = label_file.readlines()

    for idx, line in enumerate(lines):
        parts = line.strip().split()
        if len(parts) < 5:
            continue
        class_id, x_center, y_center, width, height = map(float, parts[:5])

        # 计算原始bbox的坐标
        x_center *= img_width
        y_center *= img_height
        width *= img_width
        height *= img_height

        # 扩大bbox
        new_width = width * scale_factor
        new_height = height * scale_factor

        # 计算新的坐标
        x_min = int(x_center - new_width / 2)
        y_min = int(y_center - new_height / 2)
        x_max = int(x_center + new_width / 2)
        y_max = int(y_center + new_height / 2)

        # 确保坐标不越界
        x_min = max(0, x_min)
        y_min = max(0, y_min)
        x_max = min(img_width, x_max)
        y_max = min(img_height, y_max)

        # 裁剪图片
        cropped_img = img.crop((x_min, y_min, x_max, y_max))

        # 创建输出文件夹
        os.makedirs(output_folder, exist_ok=True)
        output_image_path = os.path.join(output_folder, f"{os.path.basename(image_path).split('.')[0]}_{idx}.jpg")
        cropped_img.save(output_image_path, quality=quality)

    print(f"Cropped and saved image: {image_path} to {output_folder}")

def process_image(image_path, txt_file, output_folder, scale_factor=1.5):
    """
    处理每张图片：查找对应的标注文件并裁剪
    :param image_path: 图片路径
    :param txt_file: 包含图片路径的txt文件
    :param output_folder: 输出文件夹
    :param scale_factor: 放大倍数
    """
    # 查找对应的标签文件
    base_name = os.path.basename(image_path)
    label_path = os.path.splitext(image_path)[0] + ".txt"
    
    if not os.path.exists(label_path):
        print(f"Label file not found for {image_path}")
        return

    # 裁剪并保存
    crop_and_resize(image_path, label_path, output_folder, scale_factor)

def process_txt_file(txt_file, output_folder, scale_factor=1.5):
    """
    读取txt文件，处理每张图片
    :param txt_file: 包含图片路径的txt文件
    :param output_folder: 输出文件夹
    :param scale_factor: 放大倍数
    """
    if not os.path.exists(txt_file):
        print(f"File {txt_file} not found.")
        return

    with open(txt_file, "r") as file:
        image_paths = [line.strip() for line in file if "OD202" in line]

    # 创建线程池进行并行处理
    with ThreadPoolExecutor() as executor:
        for image_path in image_paths:
            executor.submit(process_image, image_path, txt_file, output_folder, scale_factor)


if __name__ == "__main__":
    # 输入的txt文件路径
    txt_file = "gt_SN8_20241229.txt"  # 请替换为你的txt文件路径
    output_folder = "bbox_shelter_crop"  # 输出文件夹
    scale_factor = 1.25  # 设置缩放倍数

    process_txt_file(txt_file, output_folder, scale_factor)
